import matplotlib.pyplot as plt
import seaborn as sns
import torch

class SentimentPlot(object):
    def __init__(self, sentiment_model, bins=20, xlabel="", ylabel="", title=""):
        self.s_model = sentiment_model
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        plt.xlim(0, 1)
        plt.ylim(0, 10)
        self.bins = bins

    def textseq2sentiseq(self, text_seqs):
        senti_seqs = []
        with torch.no_grad():
            for seq in text_seqs:
                senti_scores = self.s_model(seq)
                senti_seqs.append(senti_scores)
        return senti_seqs

    def senti_shift(self, senti_seqs, start=0, end=1, legend=""):
        colors = [(1, 0, 0), (1, 1, 0), (0, 1, 0,), (0, 0, 1)]
        pos_head = [senti[start] for senti in senti_seqs]
        pos_tail = [senti[end] for senti in senti_seqs]
        sns.distplot(pos_head, bins=self.bins, rug=False, kde=True, hist=True, norm_hist=True, label="source",
                     hist_kws={"histtype": "step", "linewidth": 2,
                               "alpha": 1}, kde_kws={"color": colors[0], "lw": 0, "label": ""})
        sns.distplot(pos_tail, bins=self.bins, rug=False, kde=True, hist=True, norm_hist=True, label="tail",
                     hist_kws={"histtype": "step", "linewidth": 2,
                               "alpha": 1}, kde_kws={"color": colors[0], "lw": 0, "label": ""})